////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#pragma once


#ifndef __WN_CONCURRENCY_THREAD_POOL_INL__
#define __WN_CONCURRENCY_THREAD_POOL_INL__

#ifdef _WN_WINDOWS
#include <windows.h>
#endif
#include <algorithm>

using namespace WNConcurrency;

namespace WNConcurrency {
    struct ThreadStartData {
        WNThreadPool* mThreadPool;
        WN_UINT32     mThreadNum;
    };
}

WNThreadPool::WNThreadPool() : 
    mNumThreads(0),
    mShutdown(WN_TRUE)
#ifdef _WN_WINDOWS
    ,mIOCPHandle(INVALID_HANDLE_VALUE) 
#endif 
{
}

WNThreadPool::~WNThreadPool() {
    Cleanup();
}
WN_VOID WNThreadPool::WNThreadPoolFunction(ThreadStartData* _data) {
    WNThreadPool* pool = _data->mThreadPool;
    WN_UINT32 threadNum = _data->mThreadNum;
    pool->mThreadStartMutex.Unlock();

#ifdef _WN_WINDOWS
    LPOVERLAPPED overlapped = NULL;
    DWORD bytesTransferred = 0;
    ULONG_PTR completionKey = 0;
    while(WN_TRUE) {
        BOOL ret = GetQueuedCompletionStatus(
                                    pool->mIOCPHandle,  
                                    &bytesTransferred, 
                                    &completionKey, 
                                    &overlapped, 
                                    INFINITE);
        if(0 == completionKey) {
            //Message from the pool, not an actual job
            switch(bytesTransferred) {
                case 0:
                    //Cleanup Message (We should be responsible for freeing memory
                    pool->mThreadMutex.Lock();
                        WN_DELETE(pool->mThreads[threadNum]);
                        pool->mThreadFreelist.push_back(threadNum);
                    pool->mThreadMutex.Unlock();
                    return;
                    break;
                case 1:
                    //Cleanup Message (Pool will worry about freeing memory)
                    return;
                    break;
                default:
                    WN_RELEASE_ASSERT_DESC(WN_FALSE, "This isnt right, we got a message from the queue we don't understand");
                    break;
            }
        } else { 
            WNContainers::WNResourcePointer<WNThreadJob>* job = reinterpret_cast<WNContainers::WNResourcePointer<WNThreadJob>* >(completionKey);
            (*job)->Execute();
            (*job)->mCompleted = WN_TRUE;
            WN_DELETE(job);
        }
    }
#elif defined(_WN_LINUX)
    while(WN_TRUE) {
        pool->mNumJobs.Lock();
        //If we get past this lock, it means there is at least one job available
        pool->mJobMutex.Lock();
        WNContainers::WNResourcePointer<WNThreadJob> job = pool->mJobs.front();
        pool->mJobs.pop_front();
        pool->mJobMutex.Unlock();
        if(!job) {
            //Got the signal to get out of here
            pool->mThreadMutex.Lock();
                pool->mThreadFreelist.push_back(threadNum);
            pool->mThreadMutex.Unlock();
            return;
        } else {
            job->Execute();
            job->mCompleted = WN_TRUE;
        }
    }
#endif
}

WN_VOID WNThreadPool::Cleanup() {
#ifdef _WN_WINDOWS
    mThreadMutex.Lock();
    for(WN_SIZE_T i = 0; i < mThreads.size(); ++i) {
        if(std::find(mThreadFreelist.begin(), mThreadFreelist.end(), i) == mThreadFreelist.end()) {
            PostQueuedCompletionStatus(mIOCPHandle, 1, NULL, 0);
        }
    } 
    for(WN_SIZE_T i = 0; i < mThreads.size(); ++i) {
        if(std::find(mThreadFreelist.begin(), mThreadFreelist.end(), i) == mThreadFreelist.end()) {
            mThreads[i]->Join();
            WN_DELETE(mThreads[i]);
        }
    } 
    mThreads.clear();
    mThreadFreelist.clear();
    mThreadMutex.Unlock();
    CloseHandle(mIOCPHandle);
    mIOCPHandle = INVALID_HANDLE_VALUE;
#elif defined(_WN_LINUX)
    mThreadMutex.Lock();
    mShutdown = WN_TRUE;
    for(WN_SIZE_T i = 0; i < mThreads.size(); ++i) {
        if(std::find(mThreadFreelist.begin(), mThreadFreelist.end(), i) == mThreadFreelist.end()) {
            mJobMutex.Lock();
                mJobs.push_back(WNContainers::WNResourcePointer<WNThreadJob>());
            mJobMutex.Unlock();
            mNumJobs.Unlock();
        }
    }
    mThreadMutex.Unlock();
    for(WN_SIZE_T i = 0; i < mThreads.size(); ++i) {
        if(std::find(mThreadFreelist.begin(), mThreadFreelist.end(), i) == mThreadFreelist.end()) {
            mThreads[i]->Join();
            WN_DELETE(mThreads[i]);
        }
    } 
    mThreads.clear();
    mThreadFreelist.clear();
#endif
}

WNThreadPoolError WNThreadPool::Initialize(WN_UINT32 _numWorkers) {
    if(mShutdown == WN_FALSE) {
        return(eWNReInitialized);
    }
#ifdef _WN_WINDOWS
    mIOCPHandle = CreateIoCompletionPort(INVALID_HANDLE_VALUE, WN_NULL, 0, 0);
    if(INVALID_HANDLE_VALUE == mIOCPHandle) {
        return(eWNPlatformError);
    }
#endif
    ThreadStartData dat = {0};
    dat.mThreadPool = this;
    for(WN_UINT32 i = 0; i < _numWorkers; ++i) {
        dat.mThreadNum = i;
        WNConcurrency::WNThread* thread = WN_NEW(WNConcurrency::WNThread(&WNThreadPool::WNThreadPoolFunction, &dat));
        if(!thread->IsValid()) {
            WN_DELETE(thread);
            Cleanup();
            return(eWNInvalidParameters);
        }
        mThreads.push_back(thread);
        mThreadStartMutex.Lock();
    }
    mNumThreads = _numWorkers;
    mShutdown = WN_FALSE;
    return(eWNOK);
}

WNThreadPoolError WNThreadPool::AddJob(const WNContainers::WNResourcePointer<WNThreadJob>& _job) {
#ifdef _WN_WINDOWS
    //The job is responsible for deleting this data
    //At some point when WNMemory is more completed, should probably switch this to a bucketed allocator (reduce fragmentation, and allocation/free time)
    WNContainers::WNResourcePointer<WNThreadJob>* mJob = WN_NEW(WNContainers::WNResourcePointer<WNThreadJob>(_job));
    if(!PostQueuedCompletionStatus(mIOCPHandle, 0, reinterpret_cast<ULONG_PTR>(mJob), 0))
    {
        WN_DELETE(mJob);
        return(eWNUninitialized);
    }
#elif defined _WN_LINUX
    mJobMutex.Lock();
    if(mShutdown) {
        mJobMutex.Unlock();
        return(eWNUninitialized);
    }
    mJobs.push_back(_job);
    mJobMutex.Unlock();
    mNumJobs.Unlock();
#endif
    return(eWNOK);
}

#endif // __WN_CONCURRENCY_THREAD_POOL_INL__
